#include "kernel_operator.h"
#include "my_functions.h"
#include "op_common.h"
#include <cstdint>
#include <type_traits>

namespace AscendC {

constexpr static float ERF_PARAM1 = -0.3512339572e-8;
constexpr static float ERF_PARAM2 = 0.2645266170e-6;
constexpr static float ERF_PARAM3 = -0.7929488134e-5;
constexpr static float ERF_PARAM4 = 0.1106123840e-3;
constexpr static float ERF_PARAM5 = 0.6518995814e-4;
constexpr static float ERF_PARAM6 = -0.7266616915e-1;
constexpr static float ERF_PARAM7 = -0.1595769883e1;
constexpr static float ERF_THRESHOLD = 5.75;

template <typename _DT_X, typename _DT_Y> class Kernel {
public:
  TPipe pipe;
  int64_t size = 0;
  int64_t tileLength = 0;
  int64_t finalLength = 0;
  DefInTensor(X);
  DefOutTensor(Y);
  DefBufVECCALC(TEMPB);
  DefBufVECCALC(TEMP1);

public:
  __aicore__ inline Kernel() {}

  template <typename T>
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, T tiling_data) {
    this->size = tiling_data.size;
    this->tileLength = tiling_data.tileLength;
    this->finalLength = tiling_data.finalLength;
    // 切分global
    GSetBuffer(X, x, 0, tiling_data.size);
    GSetBuffer(Y, y, 0, tiling_data.size);
    // 初始化队列
    InitQueueSimple(X, tileLength);
    InitQueueSimple(Y, tileLength);
    // 初始化BUF
    const auto bufferSize = tileLength * sizeof(_DT_X);
    InitTBufBuffer(TEMPB, bufferSize);
    if constexpr (std::is_same_v<_DT_X, float>) {
      InitTBufBuffer(TEMP1, bufferSize);
    }
  }

  __aicore__ inline void Process() {
    auto loopCount = size / tileLength;
    for (uint32_t i = 0; i < loopCount; ++i) {
      CopyIn(i, tileLength);
      Compute(tileLength);
      CopyOut(i, tileLength);
    }
    if (finalLength > 0) {
      auto copyCount = alignToBlock<_DT_X>(finalLength);
      CopyIn(loopCount, copyCount);
      Compute(finalLength);
      CopyOut(loopCount, copyCount);
    }
  }

  __aicore__ inline void CopyIn(uint32_t i, uint32_t calcCount) {
    EnQueGlobal2Local(X, i * tileLength, calcCount);
  }

  __aicore__ inline void CopyOut(uint32_t i, uint32_t calcCount) {
    DeQueLocal2Global(Y, i * tileLength, calcCount);
  }

  __aicore__ inline void Compute(uint32_t calcCount) {
    DeQueSimple(X);
    QueAllocSimple(Y);
    if constexpr (std::is_same_v<_DT_X, half>) {
      TBufGet(TEMPB, half);
      myGelu(LTensorName(Y), LTensorName(X), LTensorName(Y),
             TBufTensorName(TEMPB), calcCount);
    } else {
      // const scalar_t kAlpha = scalar_t(M_SQRT1_2);
      TBufGet(TEMPB, float);
      TBufGet(TEMP1, float);

      // Muls(LTensor(Y), LTensor(X), -1.702f, calcCount);
      // Exp(LTensor(Y), LTensor(Y), calcCount);
      // Adds(LTensor(Y), LTensor(Y), 1.0f, calcCount);
      // Mul(LTensor(Y), LTensor(Y), LTensor(X), calcCount);

      // res =
      // x/(1+exp(((((((a1*x^2+a2)*x^2+a3)*x^2+a4)*x^2+a5)*x^2+a6)*x^2+a7)*x))
      Mins(BTensor(TEMPB), LTensor(X), ERF_THRESHOLD, calcCount);

      Mul(BTensor(TEMP1), BTensor(TEMPB), BTensor(TEMPB), calcCount);
      Muls(LTensor(Y), BTensor(TEMP1), ERF_PARAM1, calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM2, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMP1), calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM3, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMP1), calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM4, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMP1), calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM5, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMP1), calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM6, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMP1), calcCount);

      Adds(LTensor(Y), LTensor(Y), ERF_PARAM7, calcCount);
      Mul(LTensor(Y), LTensor(Y), BTensor(TEMPB), calcCount);

      Exp(LTensor(Y), LTensor(Y), calcCount);

      Adds(LTensor(Y), LTensor(Y), 1.0f, calcCount);
      Div(LTensor(Y), LTensor(X), LTensor(Y), calcCount);

      // constexpr float BETA = 0.044715;
      // constexpr float ALPHA = -1.5957691;
      // // x^2
      // Mul(BTensor(TEMP1), LTensor(X), LTensor(X), calcCount);
      // // x^3
      // Mul(BTensor(TEMP1), LTensor(X), BTensor(TEMP1), calcCount);
      // // 0.044715 * x^3
      // Muls(BTensor(TEMP1), BTensor(TEMP1), BETA, calcCount);
      // // x + 0.044715 * x^3
      // Add(BTensor(TEMP1), LTensor(X), BTensor(TEMP1), calcCount);
      // // -sqrt(8/pi)(x + 0.044715 * x^3)
      // Muls(BTensor(TEMP1), BTensor(TEMP1), ALPHA, calcCount);
      // // exp(-sqrt(8/pi)(x + 0.044715 * x^3))
      // Exp(BTensor(TEMP1), BTensor(TEMP1), calcCount);
      // // 1 + exp(-sqrt(8/pi)(x + 0.044715 * x^3))
      // Adds(BTensor(TEMP1), BTensor(TEMP1), 1.0f, calcCount);
      // // x / (1 + exp(-sqrt(8/pi)(x + 0.044715 * x^3)))
      // Div(LTensor(Y), LTensor(X), BTensor(TEMP1), calcCount);

      // const float kAlpha = M_SQRT1_2;
      // // return x * scalar_t(0.5) * (scalar_t(1) + std::erf(x * kAlpha));
      // Muls(BTensor(TEMPC), LTensor(X), kAlpha, calcCount);
      // my_erf2(LTensor(Y), BTensor(TEMPC), BTensor(TEMP1), BTensor(TEMP2),
      //         BTensor(TEMP3), BTensor(TEMP4), calcCount);
      // // my_erf(LTensor(Y), BTensor(TEMPC), BTensor(CMPSUM), BTensor(CMP),
      // //        BTensor(CMP2), BTensor(CMP3), BTensor(TEMP1), BTensor(TEMP2),
      // //        BTensor(TEMP3), BTensor(TEMP4), BTensor(TEMP5), calcCount);
      // Adds(BTensor(TEMPC), LTensor(Y), 1.0f, calcCount);
      // Muls(LTensor(Y), BTensor(TEMPC), 0.5f, calcCount);
      // Mul(LTensor(Y), LTensor(Y), LTensor(X), calcCount);
      // const float kBeta = M_SQRT2 * M_2_SQRTPI * 0.5f;
      // const float kKappa = 0.044715;
      // my_tanhf(LTensor(Y), LTensor(X), BTensor(TEMP1), BTensor(TEMP2),
      //          BTensor(TEMP3), BTensor(TEMP4), BTensor(TEMP5),
      //          BTensor(TEMP6), BTensor(T0), BTensor(T1), BTensor(T2),
      //          BTensor(T22), BTensor(TInf), BTensor(TNInf),
      //          BTensor(Tl2PowerN55), calcCount);
      // // x * x * x
      // Mul(BTensor(TEMPB), LTensor(X), LTensor(X), calcCount);
      // Mul(BTensor(TEMPB), BTensor(TEMPB), LTensor(X), calcCount);
      // // kBeta * (x + kKappa * x_cube)
      // Muls(BTensor(TEMPB), BTensor(TEMPB), kKappa, calcCount);
      // Add(BTensor(TEMPB), BTensor(TEMPB), LTensor(X), calcCount);
      // Muls(BTensor(TEMPB), BTensor(TEMPB), kBeta, calcCount);
      // // 计算tanh
      // // scalar_t(0.5) * x * (scalar_t(1) + std::tanh(inner))
      // Add(BTensor(TEMPC), BTensor(TEMPC), BTensor(T1), calcCount);
      // Mul(BTensor(TEMPC), BTensor(TEMPC), LTensor(X), calcCount);
      // Muls(LTensor(Y), BTensor(TEMPC), 0.5f, calcCount);
    }
    EnQue(Y);
    QueFree(X);
  }
};
} // namespace AscendC

extern "C" __global__ __aicore__ void gelu(GM_ADDR x, GM_ADDR y,
                                           GM_ADDR workspace, GM_ADDR tiling) {
  GET_TILING_DATA(tiling_data, tiling);
  AscendC::Kernel<DTYPE_X, DTYPE_Y> op;
  op.Init(x, y, tiling_data);
  op.Process();
}